import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import time

# 修改后的希尔排序算法
def shell_sort(arr):
    n = len(arr)
    steps = []
    gap = n // 2  # 初始间隔设置为长度的一半

    while gap > 0:
        for i in range(gap, n):
            temp = arr[i]
            j = i
            while j >= gap and arr[j - gap] > temp:
                steps.append((arr.copy(), j, j - gap))  # 记录排序过程和当前比较的位置
                arr[j] = arr[j - gap]
                j -= gap
                steps.append((arr.copy(), j, j - gap))  # 记录排序过程和当前比较的位置
            steps.append((arr.copy(), j, j))
            arr[j] = temp
            steps.append((arr.copy(), j, j))  # 记录排序过程和当前比较的位置
        gap //= 2  # 缩小间隔

    return steps

# 随机生成一个数组
arr = np.random.randint(0, 100, 20)

# 获取排序步骤
steps = shell_sort(arr.copy())

# 创建动画
fig, ax = plt.subplots()
ax.set_xlim(-1, len(arr))
ax.set_ylim(0, max(arr) + 10)

bar_rects = ax.bar(range(len(arr)), arr, align="edge")

def animate(i):
    current_step = steps[i][0]
    compare_idx1 = steps[i][1]
    compare_idx2 = steps[i][2]

    for idx, (rect, val) in enumerate(zip(bar_rects, current_step)):
        if idx == compare_idx1 or idx == compare_idx2:
            rect.set_color('r')  # 将比较的柱子设为红色
            time.sleep(0.5)
        else:
            rect.set_color('b')  # 其余柱子设为蓝色
        rect.set_height(val)
    time.sleep(0.05)  # 可以适当调整延迟以观察排序过程

    return bar_rects

ani = animation.FuncAnimation(fig, animate, frames=len(steps), interval=50, blit=True)

plt.show()
